import json
import torch
from torch import nn,optim

from transformers import BertModel,BertTokenizer,set_seed,get_linear_schedule_with_warmup
import logging
import numpy as np
import time
import config

def load_json(data_path):
    with open(data_path, encoding="utf-8") as f:
        return json.loads(f.read())

# device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
device = config.device_type
# model_path =  '/home/ideal/NLP/zhangli/pre_model/roberta_ch'
model_path = config.model_path

# with open ('data/2danwei_label.json') as f:
with open (config.label2idx_path) as f:
    l = f.readline()
    name_label = json.loads(l)
label_name = {}
for key in name_label.keys():
    label_name[name_label[key]]= key
danwei_id_path = config.danwei_id_path

danwei_id = load_json(danwei_id_path)

class Bert_Model(nn.Module):
    def __init__(self,categories):
        super(Bert_Model,self).__init__()
        self.tokenizer = BertTokenizer.from_pretrained(pretrained_model_name_or_path=model_path)
        self.bert = BertModel.from_pretrained(model_path)
        self.clssifier = nn.Linear(768,categories)

    def forward(self,batch_sentences):
        sentence_tokenized = self.tokenizer(
            batch_sentences,
            truncation = True,
            padding = True,
            max_length = max_len,
            add_special_tokens = True
            )
        input_ids = torch.tensor(sentence_tokenized['input_ids']).to(device)
        attention_mask = torch.tensor(sentence_tokenized['attention_mask']).to(device)
        bert_output = self.bert(input_ids,attention_mask)      
        bert_cls = bert_output[0][:,0,:]
        #bert_cls = bert_output[1]
        linear_output = self.clssifier(bert_cls)
        return linear_output

# def inference(text_id,test_sentence,model_version):   
    # batch_size = 1
    # bert_classify = Bert_Model(category).to(device)
    # state_dict = torch.load(model_version)
    # bert_classify.load_state_dict(state_dict['model'])
def inference(test_sentence):  
    outputs = bert_classify([test_sentence])
    outputs = nn.Softmax(dim=1)(outputs)
    outputs = outputs.argmax(dim=1)
    outputs = outputs.cpu().tolist()
    # return {text_id:text_id,result:outputs}
    result = label_name[outputs[0]]
    # id = danwei_id[result]
    return result

category = len(label_name)
max_len= 256
model_version=config.model_version
# model_version = "/home/ideal/NLP/zhangli/classify/paidan/model/paidan_ai3_roberta_1.pth"
bert_classify = Bert_Model(category).to(device)
state_dict = torch.load(model_version,map_location='cpu')
# bert_classify.load_state_dict(state_dict['model'])
bert_classify.load_state_dict(state_dict['model'],strict=False)



if __name__ == '__main__':
    test_sentence =['浦东新区,科教文卫类,教育,学校管理,教师素质,【投诉】市民来电反映:投诉上述学校三（2）班的班主任、语文老师李崇，该老师每学一篇课文就划50-60个词语，第二天就默写，默的内容还有拓展，导致很多学生写不出，写错的一个词语要订正4遍，错的多的一次要订正一百多遍，市民认为这是变相的体罚；该老师还要求考卷上的字一定要在田字格中间，笔画都不能碰到田字格边框，不然写对也算写错；要求小朋友说每句话之前都要加上李老师三个字，但李崇老师对于学生家长的消息从来没有答复；整个班级的考试成绩没有别的班级好，李崇老师只会说因为男孩子不做作业，女孩子抄作业，将问题归咎在家长头上，也不告知怎么提高怎么改正；之前考试期间还因为一个学生座位周围比较脏，就不发试卷，等学生捡完垃圾去拿试卷，李崇老师却说“我让你过来了吗”，还曾对学生说“你说的话家里人相信，我不会相信的”，市民认为这根本不是一个为人师表的人可以说出来的话。小朋友们对于该老师都有怨言，认为老师很凶，日常很少鼓励，整个班级气氛压抑，只要上学就觉得头痛的厌学、抑郁倾向，大部分同学语文作业要写到晚上22:00以后，甚至23：00后。诉求：希望管理部门核实，要求立即更换老师。（市民要求信息保密 无需回复）']#输入模型必须是列表形式
    # category = 2
    # max_len= 256
    t1=time.time()
    text_id='XXX'
    # model_version="model/tufa_class_ro_3.pth"
    model_version = "train_model/paidan_ai3_roberta_1.pth"
    result = inference(test_sentence)
    result = label_name[result]
    # print(result.item())
    print(result)
    print('单条预测时长：',(time.time()-t1)*1000)


